/****************************************************************************
 * Copyright (c) 2018-2023 by the Cabana authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Cabana library. Cabana is distributed under a   *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#include <Cabana_Grid_Array.hpp>
#include <Cabana_Grid_GlobalGrid.hpp>
#include <Cabana_Grid_GlobalMesh.hpp>
#include <Cabana_Grid_HypreStructuredSolver.hpp>
#include <Cabana_Grid_Partitioner.hpp>
#include <Cabana_Grid_ReferenceStructuredSolver.hpp>
#include <Cabana_Grid_Types.hpp>

#include <Kokkos_Core.hpp>

#include <gtest/gtest.h>

#include <array>
#include <vector>

using namespace Cabana::Grid;

namespace Test
{

//---------------------------------------------------------------------------//
// FIXME: Only run test if HYPRE is compatible with the memory space. This
// is currently written in this structure because HYPRE only has
// compile-time switches for backends and hence only one can be used at a
// time. Once they have a run-time switch we can use that instead.
template <class MemorySpace>
std::enable_if_t<!HypreIsCompatibleWithMemorySpace<MemorySpace>::value, void>
poissonTest( const std::string&, const std::string&, MemorySpace )
{
}

template <class MemorySpace>
std::enable_if_t<HypreIsCompatibleWithMemorySpace<MemorySpace>::value, void>
poissonTest( const std::string& solver_type, const std::string& precond_type,
             MemorySpace )
{
    // Create the global grid.
    double cell_size = 0.2;
    std::array<bool, 2> is_dim_periodic = { false, false };
    std::array<double, 2> global_low_corner = { -1.0, -2.0 };
    std::array<double, 2> global_high_corner = { 1.0, 1.0 };
    auto global_mesh = createUniformGlobalMesh( global_low_corner,
                                                global_high_corner, cell_size );

    // Create the global grid.
    DimBlockPartitioner<2> partitioner;
    auto global_grid = createGlobalGrid( MPI_COMM_WORLD, global_mesh,
                                         is_dim_periodic, partitioner );

    // Create a local grid.
    auto local_mesh = createLocalGrid( global_grid, 1 );
    auto owned_space = local_mesh->indexSpace( Own(), Cell(), Local() );

    // Create the RHS.
    auto vector_layout = createArrayLayout( local_mesh, 1, Cell() );
    auto rhs = createArray<double, MemorySpace>( "rhs", vector_layout );
    ArrayOp::assign( *rhs, 1.0, Own() );

    // Create the LHS.
    auto lhs = createArray<double, MemorySpace>( "lhs", vector_layout );
    ArrayOp::assign( *lhs, 0.0, Own() );

    HYPRE_Init();
    {
        // Create a solver.
        auto solver = createHypreStructuredSolver<double, MemorySpace>(
            solver_type, *vector_layout );

        // Create a 5-point 2d laplacian stencil.
        std::vector<std::array<int, 2>> stencil = {
            { 0, 0 }, { -1, 0 }, { 1, 0 }, { 0, -1 }, { 0, 1 } };
        solver->setMatrixStencil( stencil );

        // Create the matrix entries. The stencil is defined over cells.
        auto matrix_entry_layout = createArrayLayout( local_mesh, 5, Cell() );
        auto matrix_entries = createArray<double, MemorySpace>(
            "matrix_entries", matrix_entry_layout );
        auto entry_view = matrix_entries->view();
        Kokkos::parallel_for(
            "fill_matrix_entries",
            createExecutionPolicy( owned_space, TEST_EXECSPACE() ),
            KOKKOS_LAMBDA( const int i, const int j ) {
                entry_view( i, j, 0 ) = 4.0;
                entry_view( i, j, 1 ) = -1.0;
                entry_view( i, j, 2 ) = -1.0;
                entry_view( i, j, 3 ) = -1.0;
                entry_view( i, j, 4 ) = -1.0;
            } );

        solver->setMatrixValues( *matrix_entries );

        // Set the tolerance.
        solver->setTolerance( 1.0e-9 );

        // Set the maximum iterations.
        solver->setMaxIter( 2000 );

        // Set the print level.
        solver->setPrintLevel( 1 );

        // Create a preconditioner.
        if ( "none" != precond_type )
        {
            auto preconditioner =
                createHypreStructuredSolver<double, MemorySpace>(
                    precond_type, *vector_layout, true );
            solver->setPreconditioner( preconditioner );
        }

        // Setup the problem.
        solver->setup();

        // Solve the problem.
        solver->solve( *rhs, *lhs );

        // Create a solver reference for comparison.
        auto lhs_ref =
            createArray<double, MemorySpace>( "lhs_ref", vector_layout );
        ArrayOp::assign( *lhs_ref, 0.0, Own() );

        auto ref_solver = createReferenceConjugateGradient<double, MemorySpace>(
            *vector_layout );
        ref_solver->setMatrixStencil( stencil );
        const auto& ref_entries = ref_solver->getMatrixValues();
        auto matrix_view = ref_entries.view();
        auto global_space = local_mesh->indexSpace( Own(), Cell(), Global() );
        int ncell_i = global_grid->globalNumEntity( Cell(), Dim::I );
        int ncell_j = global_grid->globalNumEntity( Cell(), Dim::J );
        Kokkos::parallel_for(
            "fill_ref_entries",
            createExecutionPolicy( owned_space, TEST_EXECSPACE() ),
            KOKKOS_LAMBDA( const int i, const int j ) {
                int gi =
                    i + global_space.min( Dim::I ) - owned_space.min( Dim::I );
                int gj =
                    j + global_space.min( Dim::J ) - owned_space.min( Dim::J );
                matrix_view( i, j, 0 ) = 4.0;
                matrix_view( i, j, 1 ) = ( gi > 0 ) ? -1.0 : 0.0;
                matrix_view( i, j, 2 ) = ( gi < ncell_i - 1 ) ? -1.0 : 0.0;
                matrix_view( i, j, 3 ) = ( gj > 0 ) ? -1.0 : 0.0;
                matrix_view( i, j, 4 ) = ( gj < ncell_j - 1 ) ? -1.0 : 0.0;
            } );

        std::vector<std::array<int, 2>> diag_stencil = { { 0, 0 } };
        ref_solver->setPreconditionerStencil( diag_stencil );
        const auto& preconditioner_entries =
            ref_solver->getPreconditionerValues();
        auto preconditioner_view = preconditioner_entries.view();
        Kokkos::parallel_for(
            "fill_preconditioner_entries",
            createExecutionPolicy( owned_space, TEST_EXECSPACE() ),
            KOKKOS_LAMBDA( const int i, const int j ) {
                preconditioner_view( i, j, 0 ) = 1.0 / 4.0;
            } );

        ref_solver->setTolerance( 1.0e-11 );
        ref_solver->setPrintLevel( 1 );
        ref_solver->setup();
        ref_solver->solve( *rhs, *lhs_ref );

        // Check the results.
        auto lhs_host = Kokkos::create_mirror_view_and_copy(
            Kokkos::HostSpace(), lhs->view() );
        auto lhs_ref_host = Kokkos::create_mirror_view_and_copy(
            Kokkos::HostSpace(), lhs_ref->view() );
        for ( int i = owned_space.min( Dim::I ); i < owned_space.max( Dim::I );
              ++i )
            for ( int j = owned_space.min( Dim::J );
                  j < owned_space.max( Dim::J ); ++j )
                EXPECT_FLOAT_EQ( lhs_host( i, j, 0 ), lhs_ref_host( i, j, 0 ) );

        // Setup the problem again. We would need to do this if we changed the
        // matrix entries.
        solver->setup();

        // Solve the problem again
        ArrayOp::assign( *rhs, 2.0, Own() );
        ArrayOp::assign( *lhs, 0.0, Own() );
        solver->solve( *rhs, *lhs );

        // Compute another reference solution.
        ArrayOp::assign( *lhs_ref, 0.0, Own() );
        ref_solver->solve( *rhs, *lhs_ref );

        // Check the results again
        lhs_host = Kokkos::create_mirror_view_and_copy( Kokkos::HostSpace(),
                                                        lhs->view() );
        lhs_ref_host = Kokkos::create_mirror_view_and_copy( Kokkos::HostSpace(),
                                                            lhs_ref->view() );
        for ( int i = owned_space.min( Dim::I ); i < owned_space.max( Dim::I );
              ++i )
            for ( int j = owned_space.min( Dim::J );
                  j < owned_space.max( Dim::J ); ++j )
                EXPECT_FLOAT_EQ( lhs_host( i, j, 0 ), lhs_ref_host( i, j, 0 ) );
    }
    HYPRE_Finalize();
}

//---------------------------------------------------------------------------//
// RUN TESTS
//---------------------------------------------------------------------------//
TEST( HypreStructuredSolver, PCGnone2d )
{
    poissonTest( "PCG", "none", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, GMRESnone2d )
{
    poissonTest( "GMRES", "none", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, BiCGSTABnone2d )
{
    poissonTest( "BiCGSTAB", "none", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, PFMGnone2d )
{
    poissonTest( "PFMG", "none", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, PCGdiag2d )
{
    poissonTest( "PCG", "Diagonal", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, GMRESdiag2d )
{
    poissonTest( "GMRES", "Diagonal", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, BiCGSTABdiag2d )
{
    poissonTest( "BiCGSTAB", "Diagonal", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, PCGjacobi2d )
{
    poissonTest( "PCG", "Jacobi", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, GMRESjacobi2d )
{
    poissonTest( "GMRES", "Jacobi", TEST_MEMSPACE{} );
}

TEST( HypreStructuredSolver, BiCGSTABjacobi2d )
{
    poissonTest( "BiCGSTAB", "Jacobi", TEST_MEMSPACE{} );
}

//---------------------------------------------------------------------------//

} // end namespace Test
